{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from src import const\n",
    "import pandas as pd\n",
    "import skimage\n",
    "from skimage import io\n",
    "import numpy as np     \n",
    "from src.dataset import CelebaAttDataset\n",
    "import torch\n",
    "from src.networks import ModuleWithAttr\n",
    "from torch import nn\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('./data/info.csv', index_col='image_name')\n",
    "train_df = df[df.status == 'train']\n",
    "val_df = df[df.status == 'val']\n",
    "test_df = df[df.status == 'test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = CelebaAttDataset(train_df, mode='RESCALE192_160')\n",
    "train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=const.BATCH_SIZE, shuffle=True, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = iter(train_dataloader).next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 3, 192, 160])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample['image'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Conv_BN_Pooling(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_channels, out_channels, conv_kernel, conv_padding, pooling_kernel=None):\n",
    "        super(Conv_BN_Pooling, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, conv_kernel, stride=1, padding=conv_padding)\n",
    "        self.bn = nn.BatchNorm2d(out_channels)\n",
    "        if pooling_kernel is not None:\n",
    "            self.has_pooling = True\n",
    "            self.pooling = nn.MaxPool2d(pooling_kernel)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.bn(x)\n",
    "        if self.has_pooling:\n",
    "            x = self.pooling(x)\n",
    "        return x\n",
    "\n",
    "def func_partially_shared(SNet, TSNets, SFeature, TSFeatures):\n",
    "    '''\n",
    "    SNet: 单独的Module\n",
    "    TSNets：TSNet的list或ModuleList(所有net，单独一层)\n",
    "    SFeature：上一层SNet的输出\n",
    "    TSFeatures：上一层TSNet的输出的list\n",
    "    '''\n",
    "    assert len(TSFeatures) == len(TSNets)\n",
    "    Out_TSFeatures = []\n",
    "    for i in range(len(TSFeatures)):\n",
    "        Out_TSFeatures.append(\n",
    "            TSNets[i](\n",
    "                torch.cat([SFeature, TSFeatures[i]], dim=1)\n",
    "            )\n",
    "        )\n",
    "    Out_SFeature = SNet(\n",
    "        torch.cat([SFeature] + TSFeatures, dim=1)\n",
    "    )\n",
    "    return Out_SFeature, Out_TSFeatures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class _(object):\n",
    "    pass\n",
    "self = _()\n",
    "\n",
    "def find_idx(list_, name):\n",
    "    for i in range(len(list_)):\n",
    "        if list_[i] == name:\n",
    "            return i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MCNN_PS(ModuleWithAttr):\n",
    "    def __init__(self, label_name):\n",
    "        super(MCNN_PS, self).__init__()\n",
    "        self.label_name = label_name\n",
    "        \n",
    "        self.C_TS = [32, 64, 128, 256, 128]\n",
    "        self.C_S = [8, 16, 32, 64, 32]\n",
    "        self.groups = [\n",
    "            ['Arched_Eyebrows', 'Bags_Under_Eyes','Bald', 'Bangs', 'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Bushy_Eyebrows',\n",
    "            'Eyeglasses', 'Gray_Hair', 'Narrow_Eyes', 'Receding_Hairline', 'Wearing_Hat'],\n",
    "            ['Big_Nose', 'High_Cheekbones', 'Pointy_Nose', 'Rosy_Cheeks', 'Sideburns', 'Wearing_Earrings'],\n",
    "            ['Big_Lips', 'Double_Chin', 'Goatee', 'Mustache', \n",
    "             'Mouth_Slightly_Open', 'No_Beard', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie'],\n",
    "            ['5_o_Clock_Shadow', 'Attractive', 'Blurry', 'Chubby', 'Heavy_Makeup', 'Male', 'Oval_Face', 'Pale_Skin', \n",
    "             'Straight_Hair', 'Smiling', 'Wavy_Hair', 'Young'],\n",
    "        ]\n",
    "        self.label_name = label_name\n",
    "        self.g2num = [[find_idx(self.label_name, name) for name in group] for group in self.groups]\n",
    "\n",
    "\n",
    "        # 第一层\n",
    "        S_List = [Conv_BN_Pooling(3, self.C_S[0], 3, 1, 2)]\n",
    "        # 后面的几层\n",
    "        for i in range(1, len(self.C_S)):\n",
    "            S_List.append(Conv_BN_Pooling(self.C_S[i-1] + self.C_TS[i-1] * 4, self.C_S[i], 3, 1, 2))\n",
    "\n",
    "        TS_List = []\n",
    "        for i in range(4):\n",
    "            tmp = []\n",
    "            tmp.append(Conv_BN_Pooling(3, self.C_TS[0], 3, 1, 2))\n",
    "            for j in range(1, len(self.C_TS)):\n",
    "                tmp.append(Conv_BN_Pooling(self.C_S[j-1] + self.C_TS[j-1], self.C_TS[j], 3, 1, 2))\n",
    "            TS_List.append(nn.ModuleList(tmp))\n",
    "\n",
    "        FC_List = []\n",
    "        for i in range(4):\n",
    "            FC_List.append(\n",
    "                nn.Sequential(\n",
    "                    nn.Linear(5*6*128, 512),\n",
    "                    nn.BatchNorm1d(512),\n",
    "                    nn.Linear(512, 512),\n",
    "                    nn.BatchNorm1d(512),\n",
    "                    nn.Linear(512, len(self.groups[i]) * 2)\n",
    "                )\n",
    "            )\n",
    "\n",
    "        self.S_List = nn.ModuleList(S_List)\n",
    "        self.TS_List = nn.ModuleList(TS_List)\n",
    "        self.FC_List = nn.ModuleList(FC_List)\n",
    "        \n",
    "        self.label_loss = nn.CrossEntropyLoss()\n",
    "\n",
    "    \n",
    "    def forward(self, sample):\n",
    "        ## forward\n",
    "        # 第一层\n",
    "        Out_SFeature = self.S_List[0](sample['image'])\n",
    "        Out_TSFeatures = [self.TS_List[i][0](sample['image']) for i in range(4)]\n",
    "        ## 后面几层\n",
    "        for i in range(1, len(self.S_List)):\n",
    "            Out_SFeature, Out_TSFeatures = func_partially_shared(\n",
    "                self.S_List[i], \n",
    "                [self.TS_List[layer][i] for layer in range(4)],\n",
    "                Out_SFeature,\n",
    "                Out_TSFeatures,\n",
    "            )\n",
    "\n",
    "\n",
    "        output = [None] * len(self.label_name)\n",
    "        for i in range(4):\n",
    "            group_output = Out_TSFeatures[i].reshape(Out_TSFeatures[i].shape[0], -1)\n",
    "            group_output = self.FC_List[i](group_output)\n",
    "            group_output = group_output.reshape(group_output.size()[0], 2, -1)\n",
    "            group_output = torch.split(group_output, 1, 2)\n",
    "            # 把group的内容\n",
    "            for j, value in enumerate(group_output):\n",
    "                output[self.g2num[i][j]] = value\n",
    "        output = torch.cat(output, dim=2)\n",
    "        return output\n",
    "    \n",
    "    def cal_loss(self, sample, output):\n",
    "        loss = {\n",
    "            'all': self.label_loss(output, sample['label'])\n",
    "        }\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = MCNN_PS(label_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = net(sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'all': tensor(0.7190)}"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.cal_loss(sample, output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 13])"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "group_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 32, 6, 5])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Out_SFeature.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 128, 6, 5])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Out_TSFeatures[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
